﻿#include <QtDebug>
#include "QlyAVFrame.h"
#include "QlySwsContext.h"

//extern "C" {
//#include "libswscale/swscale.h"
//}


static enum AVPixelFormat toAVPixelFormat(QImage::Format format)
{
    switch (format) {
    case QImage::Format_Invalid:
        return AV_PIX_FMT_NONE;
    case QImage::Format_Mono:
        return AV_PIX_FMT_MONOBLACK;
    case QImage::Format_MonoLSB:
        return AV_PIX_FMT_NONE;
    case QImage::Format_Indexed8:
        return AV_PIX_FMT_PAL8;
    case QImage::Format_Alpha8:
    case QImage::Format_Grayscale8:
        return AV_PIX_FMT_GRAY8;
    case QImage::Format_Grayscale16:
        return AV_PIX_FMT_GRAY16LE;
    case QImage::Format_RGB32:
    case QImage::Format_ARGB32:
    case QImage::Format_ARGB32_Premultiplied:
        //return AV_PIX_FMT_BGRA;
        return AV_PIX_FMT_BGR0;
    case QImage::Format_RGB16:
    case QImage::Format_ARGB8565_Premultiplied:
        return AV_PIX_FMT_RGB565LE;
    case QImage::Format_RGB666:
    case QImage::Format_ARGB6666_Premultiplied:
        return AV_PIX_FMT_NONE;
    case QImage::Format_RGB555:
    case QImage::Format_ARGB8555_Premultiplied:
        return AV_PIX_FMT_BGR555LE;
    case QImage::Format_RGB888:
        return AV_PIX_FMT_RGB24;
    case QImage::Format_RGB444:
    case QImage::Format_ARGB4444_Premultiplied:
        return AV_PIX_FMT_RGB444LE;
    case QImage::Format_RGBX8888:
    case QImage::Format_RGBA8888:
    case QImage::Format_RGBA8888_Premultiplied:
        return AV_PIX_FMT_RGBA;
    case QImage::Format_BGR30:
    case QImage::Format_A2BGR30_Premultiplied:
    case QImage::Format_RGB30:
    case QImage::Format_A2RGB30_Premultiplied:
        return AV_PIX_FMT_NONE;
    case QImage::Format_RGBX64:
    case QImage::Format_RGBA64:
    case QImage::Format_RGBA64_Premultiplied:
        return AV_PIX_FMT_RGBA64LE;
    case QImage::Format_BGR888:
        return AV_PIX_FMT_BGR24;
    default:
        return AV_PIX_FMT_NONE;
    }
    return AV_PIX_FMT_NONE;
}

QlyAVFrame::QlyAVFrame()
{
    pFrame = av_frame_alloc();
}

bool QlyAVFrame::getBuffer(AVPixelFormat format, int width, int height)
{
    pFrame->width = width;
    pFrame->height = height;
    pFrame->format = format;
    if( av_frame_get_buffer(pFrame, 0) < 0 )
    {
        qWarning() << "av_frame_get_buffer() failed";
        return false;
    }
    return true;
}

QlyAVFrame::QlyAVFrame(AVFrame *p)
{
    pFrame = p;
}

QlyAVFrame::~QlyAVFrame()
{
    //av_frame_unref(pFrame);
    av_frame_free(&pFrame);
}

AVPixelFormat QlyAVFrame::format() const
{
    if(pFrame)
    {
        return (AVPixelFormat)pFrame->format;
    }
    return AV_PIX_FMT_NB;
}

QlyAVFrame::QlyAVFrame(const QlyAVFrame &source)
{
    if(pFrame != nullptr)
    {
        av_frame_free(&pFrame);
    }
    pFrame = av_frame_clone(source.pFrame);
}

QlyAVFrame::QlyAVFrame(int width, int height, AVPixelFormat format)
{
    pFrame = av_frame_alloc();
    pFrame->width = width;
    pFrame->height = height;
    pFrame->format = format;
    av_frame_get_buffer(pFrame, 0);
}

QlyAVFrame & QlyAVFrame::operator= (const QlyAVFrame & source)
{
    if(pFrame != nullptr)
    {
        av_frame_free(&pFrame);
    }
    pFrame = av_frame_clone(source.pFrame);
    return *this;
}

bool QlyAVFrame::toQImage(QImage &image)
{
    if(pFrame == nullptr) return false;
    if(image.isNull() || image.width() != pFrame->width || image.height() != pFrame->height)
    {
        image = QImage(pFrame->width, pFrame->height, QImage::Format_RGB32);
    }
    AVPixelFormat format = toAVPixelFormat(image.format());
    SwsContext *convetContext = sws_getContext(pFrame->width, pFrame->height, (AVPixelFormat)pFrame->format,
                                               pFrame->width, pFrame->height, format,
                                               SWS_POINT, nullptr, nullptr, nullptr);
    if(convetContext)
    {
        uint8_t *out_data[1];
        int out_linesize[1];

        out_data[0] = image.bits();
        out_linesize[0] = image.bytesPerLine();

        sws_scale(convetContext,
                pFrame->data, pFrame->linesize,
                0, pFrame->height,
                out_data, out_linesize);

        sws_freeContext(convetContext);
        return true;
    }
    return false;
}

QImage QlyAVFrame::toQImage(QImage::Format format)
{
    if(pFrame == nullptr)
    {
        return QImage();
    }
    SwsContext *convetContext = sws_getContext(pFrame->width, pFrame->height, (AVPixelFormat)pFrame->format,
                                               pFrame->width, pFrame->height, toAVPixelFormat(format),
                                               SWS_POINT, nullptr, nullptr, nullptr);
    if(convetContext)
    {
        QImage image(pFrame->width, pFrame->height, format);

        uint8_t *out_data[1];
        int out_linesize[1];

        out_data[0] = image.bits();
        out_linesize[0] = image.bytesPerLine();

        sws_scale(convetContext,
                pFrame->data, pFrame->linesize,
                0, pFrame->height,
                out_data, out_linesize);

        sws_freeContext(convetContext);
        return image;
    }
    return QImage();
}

bool QlyAVFrame::buildFromQVideoFrame(QVideoFrame &videoFrame, AVPixelFormat format, int pts)
{
    if(!videoFrame.isValid()) return false;

    if(pFrame == nullptr) pFrame = av_frame_alloc();
    if(pFrame == nullptr) return false;

    int width = videoFrame.width();
    int height = videoFrame.height();
    AVPixelFormat imgFmt = toAVPixelFormat(videoFrame.pixelFormat());
    if(format != pFrame->format
            || width != pFrame->width
            || height != pFrame->height)
    {
        reset();
        pFrame->format = format;
        pFrame->width = width;
        pFrame->height = height;

        if( av_frame_get_buffer(pFrame, 0) < 0 )
        {
            qWarning() << "av_frame_get_buffer() failed";
            return false;
        }

        if( av_frame_make_writable(pFrame) < 0 )
        {
            qWarning() << "av_frame_make_writable() failed";
            return false;
        }
    }
    SwsContext * pContext = sws_getContext(width, height, imgFmt,
                                           width, height, (AVPixelFormat)pFrame->format, SWS_POINT, nullptr, nullptr, nullptr);
    if(!pContext) return false;

    uint8_t *in_data[1];
    int in_linesize[1];
    videoFrame.map(QAbstractVideoBuffer::ReadOnly);//
    in_data[0] = (uint8_t *) videoFrame.bits();
    in_linesize[0] = videoFrame.bytesPerLine();

    sws_scale(pContext, in_data, in_linesize, 0, height,
              pFrame->data, pFrame->linesize);
    sws_freeContext(pContext);

    pFrame->pts = pts;
    return true;
}

bool QlyAVFrame::buildFromQImage(const QImage &image, AVPixelFormat frameFmt, int pts)
{
    if(image.isNull()) return false;

    if(pFrame == nullptr) pFrame = av_frame_alloc();
    if(pFrame == nullptr) return false;

    int width = image.width();
    int height = image.height();
    AVPixelFormat imgFmt = toAVPixelFormat(image.format());
    if(frameFmt != pFrame->format
            || width != pFrame->width
            || height != pFrame->height)
    {
        reset();
        pFrame->format = frameFmt;
        pFrame->width = width;
        pFrame->height = height;

        if( av_frame_get_buffer(pFrame, 0) < 0 )
        {
            qWarning() << "av_frame_get_buffer() failed";
            return false;
        }

        if( av_frame_make_writable(pFrame) < 0 )
        {
            qWarning() << "av_frame_make_writable() failed";
            return false;
        }
    }
    SwsContext * pContext = sws_getContext(width, height, imgFmt,
                                           width, height, (AVPixelFormat)pFrame->format, SWS_POINT, nullptr, nullptr, nullptr);
    if(!pContext) return false;

    const uint8_t *in_data[1];
    int in_linesize[1];

    in_data[0] = image.bits();
    in_linesize[0] = image.bytesPerLine();

    sws_scale(pContext, in_data, in_linesize, 0, height,
              pFrame->data, pFrame->linesize);
    sws_freeContext(pContext);

    pFrame->pts = pts;
    return true;
}

void QlyAVFrame::reset()
{
    if(pFrame) av_frame_unref(pFrame);
}
